package njust.csc.aecore.analysis

import com.android.apksig.ApkVerifier
import com.android.apksig.apk.MinSdkVersionException
import io.github.oshai.kotlinlogging.KotlinLogging
import njust.csc.aecore.entity.APK
import njust.csc.aecore.extension.compileSdkVersion
import njust.csc.aecore.extension.applicationName
import njust.csc.aecore.extension.icon
import njust.csc.aecore.utils.*
import soot.jimple.infoflow.android.manifest.ProcessManifest
import java.io.File
import java.io.FileNotFoundException
import java.nio.file.Path
import java.security.MessageDigest
import java.security.cert.X509Certificate
import kotlin.io.path.*
import kotlin.system.exitProcess

private val logger = KotlinLogging.logger {}

/**
 * @param relativePath 相对于 INPUT 的路径
 * */
class StaticExplorer(private val relativePath: String) {
    private val apk = APK()
    private var possibleIconPath = ""  // 从 manifest 中提取出来的 icon 路径，但 jadx 不一定有

    /**
     * 静态分析 apk
     * */
    init {
        val startTime = System.currentTimeMillis()
        apk.size = (PathUtils.inputResolve(relativePath).fileSize() / 1024).toInt()
        apk.absolutePath = PathUtils.inputResolve(relativePath).toAbsolutePath().toString()
        apk.relativePath = relativePath

        logger.info { "$relativePath: 收集基础信息" }
        parseManifest()

        logger.info { "$relativePath: 收集签名信息" }
        collectSignature()

        logger.info { "$relativePath: 提取 url 和 sdk 信息" }
        extractURLAndSDK()

        val endTime = System.currentTimeMillis()
        val analysisTime = (endTime - startTime) / 1000.0
        apk.time = analysisTime
    }

    /**
     * 导出 apk 到指定文件中
     * */
    fun dump(file: Path, append: Boolean = false): String {
        if (!file.parent.exists()) file.parent.createDirectories()
        val apkInfo = SerializeUtils.dump(apk, file, append)
        logger.info { "APK 信息导出至: ${file.toAbsolutePath()}" }
        return apkInfo
    }

    /**
     * 借用 FlowDroid 中的 ProcessManifest 类进行解析
     * */
    private fun parseManifest() {
        val androidManifest: ProcessManifest
        try {
            androidManifest = ProcessManifest(apk.absolutePath)
        } catch (e: RuntimeException) {
            e.printStackTrace()
            logger.info { "$relativePath: AndroidManifest 解析失败" }
            return
        } catch (e: FileNotFoundException) {
            logger.info { "$relativePath: 未找到文件" }
            return
        }
        apk.packageName = androidManifest.packageName
        apk.versionInfo.versionCode = androidManifest.versionCode
        apk.versionInfo.versionName = androidManifest.versionName
        apk.versionInfo.minSdkVersion = androidManifest.minSdkVersion
        apk.versionInfo.targetSdkVersion = androidManifest.targetSdkVersion
        apk.permissions = androidManifest.permissions
        apk.versionInfo.compileSdkVersion = androidManifest.compileSdkVersion()
        apk.name = androidManifest.applicationName()
        possibleIconPath = androidManifest.icon()
        androidManifest.close()
    }

    /**
     * 基于 apksig 库来解析 APK 的签名信息，包括 v1 v2 v3
     * */
    private fun collectSignature() {
        val verifier = ApkVerifier.Builder(File(apk.absolutePath)).build()
        val result: ApkVerifier.Result
        try {
            result = verifier.verify()
        } catch (_: MinSdkVersionException) {
            return
        }
        if (result.v1SchemeSigners.isNotEmpty()) {
            if (result.v1SchemeSigners.size != 1) {
                // 找到存在多个 signer 的 apk
                println(apk.absolutePath)
                exitProcess(-1)
            }
            for (signer in result.v1SchemeSigners) {
                val cert = signer.certificate ?: continue
                apk.v1SignatureInfo.name = signer.name
                apk.v1SignatureInfo.signatureFileName = signer.signatureFileName
                if (signer.certificate.type == "X.509") {
                    apk.v1SignatureInfo.certification.subject = cert.subjectX500Principal.name
                    apk.v1SignatureInfo.certification.signatureType = cert.sigAlgName
                    apk.v1SignatureInfo.certification.md5 = getThumbPrint(cert, "MD5")
                    apk.v1SignatureInfo.certification.sha1 = getThumbPrint(cert, "SHA-1")
                    apk.v1SignatureInfo.certification.sha256 = getThumbPrint(cert, "SHA-256")
                }
            }
        }
        if (result.v2SchemeSigners.isNotEmpty()) {
            for (signer in result.v2SchemeSigners) {
                val cert = signer.certificate ?: continue
                apk.v2SignatureInfo.index = signer.index
                if (signer.certificate.type == "X.509") {
                    apk.v2SignatureInfo.certification.subject = cert.subjectX500Principal.name
                    apk.v2SignatureInfo.certification.signatureType = cert.sigAlgName
                    apk.v2SignatureInfo.certification.md5 = getThumbPrint(cert, "MD5")
                    apk.v2SignatureInfo.certification.sha1 = getThumbPrint(cert, "SHA-1")
                    apk.v2SignatureInfo.certification.sha256 = getThumbPrint(cert, "SHA-256")
                }
            }
        }
        if (result.v3SchemeSigners.isNotEmpty()) {
            for (signer in result.v3SchemeSigners) {
                val cert = signer.certificate ?: continue
                apk.v2SignatureInfo.index = signer.index
                if (signer.certificate.type == "X.509") {
                    apk.v2SignatureInfo.certification.subject = cert.subjectX500Principal.name
                    apk.v2SignatureInfo.certification.signatureType = cert.sigAlgName
                    apk.v2SignatureInfo.certification.md5 = getThumbPrint(cert, "MD5")
                    apk.v2SignatureInfo.certification.sha1 = getThumbPrint(cert, "SHA-1")
                    apk.v2SignatureInfo.certification.sha256 = getThumbPrint(cert, "SHA-256")
                }
            }
        }
    }

    private fun getThumbPrint(cert: X509Certificate, type: String?): String {
        val md = MessageDigest.getInstance(type)
        val der = cert.encoded
        md.update(der)
        val digest = md.digest()
        return hexify(digest)
    }

    private fun hexify(bytes: ByteArray): String {
        val hexDigits = charArrayOf('0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'A', 'B', 'C', 'D', 'E', 'F')
        val buf = StringBuilder(bytes.size * 3)
        for (aByte in bytes) {
            buf.append(hexDigits[aByte.toInt() and 0xf0 shr 4])
            buf.append(hexDigits[aByte.toInt() and 0x0f])
            buf.append(' ')
        }
        return buf.toString()
    }

    /**
     * 反汇编后通过正则表达式提取 url 和 ip
     * */
    private fun extractURLAndSDK() {
        // jadx 方案
        val reversePath = PathUtils.getReversePath(relativePath)
        File(apk.absolutePath).reverseTo(reversePath)

        // 尝试得到 icon 信息
        apk.icon = extractIcon(reversePath.resolve("resources"))

        // 开始静态分析 sdk
        var sourcesDir = reversePath.resolve("sources")
        val sdks = mutableListOf<String>()
        extractSDK(sourcesDir, emptyList(), sdks)
        apk.sdks.addAll(FilterUtils.filterSdk(sdks))

        // 根据包名找到 app 源码目录
        val results: MutableList<String> = mutableListOf()
        if (apk.packageName != "") {
            apk.packageName.split(".").forEach {
                sourcesDir = sourcesDir.resolve(it)
            }
            RegexUtils.searchFileRecv(sourcesDir, results, RegexUtils.urlMatch)
            apk.urls = FilterUtils.filterUrl(results.toList()).toMutableSet()
        }
        PathUtils.deleteDir(reversePath.parent)
    }

    /**
     * 根据 jadx 反编译结果的 source 目录寻找用到的 SDK
     * @param sourceDir 源码目录
     * @param sdkName 当前 sdk 名称
     * */
    private fun extractSDK(sourceDir: Path, sdkName: List<String>, sdks: MutableList<String>): Boolean {
        if (sdkName.size == 3) {
            val packageName = sdkName.joinToString(separator = ".")
            if (!packageName.startsWith(apk.packageName)) sdks.add(packageName)
            return true
        }
        var flag = false
        for (file in sourceDir.listDirectoryEntries()) {
            if (file.isDirectory()) {
                flag = extractSDK(file, sdkName + listOf(file.name), sdks)
            }
        }
        if (sdkName.size == 2 && !flag) {
            val packageName = sdkName.joinToString(separator = ".")
            if (!packageName.startsWith(apk.packageName)) sdks.add(packageName)
            return true
        }
        return flag
    }

    /**
     * 尝试得到 icon 的 base64 编码字符串
     * @param resourceDir jadx 反编译后的 resource 目录
     * */
    private fun extractIcon(resourceDir: Path): String {
        if (possibleIconPath == "") {
            logger.info { "未找到 icon" }
            return ""
        }
        val possibleIcon = resourceDir.resolve(possibleIconPath)
        if (possibleIcon.exists()) {
            logger.info { "读取到的 icon 路径: $possibleIcon" }
            return ImageUtils.base64Encode(possibleIcon)
        } else {
            // 尝试在固定目录中寻找 icon
            val possibleDirName = listOf(
                "mipmap-hdpi", "mipmap-mdpi", "mipmap-xhdpi", "mipmap-xxhdpi", "mipmap-xxxhdpi",
                possibleIcon.parent.name.dropLast(3)
            )
            for (name in possibleDirName) {
                val possibleIcon2 = resourceDir.resolve("res").resolve(name).resolve(possibleIcon.name)
                if (possibleIcon2.exists()) {
                    return ImageUtils.base64Encode(possibleIcon2)
                }
            }
        }
        logger.info { "未找到 icon: $possibleIconPath" }
        return ""
    }

    /**
     * 获取应用包名
     * @return String 应用包名
     */
    fun getPackageName(): String {
        return apk.packageName
    }
}
